#include "FaceRecognizer.h"
#include <iostream>
#include <dlib/image_processing/render_face_detections.h>

FaceRecognizer::FaceRecognizer(ConfigParser& config) {
    // 从配置中获取阈值
    std::string threshold_str = config.get("threshold");
    threshold_ = threshold_str.empty() ? 0.6 : std::stod(threshold_str);
    
    // 获取模型路径
    shape_model_path_ = config.get("shape_model");
    resnet_model_path_ = config.get("resnet_model");
    face_lib_path_ = config.get("face_library");
    
    // 初始化人脸检测器
    detector_ = dlib::get_frontal_face_detector();
    
    // 加载特征点检测模型
    if (!shape_model_path_.empty()) {
        dlib::deserialize(shape_model_path_) >> sp_;
    } else {
        throw std::runtime_error("shape_model路径未配置");
    }
    
    // 加载人脸识别模型
    if (!resnet_model_path_.empty()) {
        dlib::deserialize(resnet_model_path_) >> net_;
    } else {
        throw std::runtime_error("resnet_model路径未配置");
    }
    
    // 加载人脸库
    if (!loadFaceLib()) {
        throw std::runtime_error("人脸库加载失败");
    }
}

bool FaceRecognizer::loadFaceLib() {
    std::ifstream file(face_lib_path_);
    if (!file) {
        std::cerr << "无法打开人脸库文件: " << face_lib_path_ << std::endl;
        return false;
    }
    
    std::string line;
    while (std::getline(file, line)) {
        std::istringstream iss(line);
        FaceRecord record;
        std::string value;
        
        // 读取人名
        if (!std::getline(iss, record.name, ',')) {
            continue; // 读取失败，跳过
        }
        
        // 读取特征向量
        dlib::matrix<float,0,1> descriptor(128);
        for (long i = 0; i < 128; ++i) {
            if (!std::getline(iss, value, ',')) {
                break; // 读取失败，跳出循环
            }
            descriptor(i) = std::stof(value);
        }
        record.descriptor = descriptor;
        face_lib_.push_back(record);
    }
    
    return !face_lib_.empty();
}

void FaceRecognizer::printFaceLibInfo() const {
    std::cout << "人脸库加载成功, 人脸数量: " << face_lib_.size() << std::endl;
    std::cout << "threshold: " << threshold_ << std::endl;
    std::cout << "人脸库信息:" << std::endl;
    for (const auto& record : face_lib_) {
        std::cout << "姓名: " << record.name 
                  << ", 特征向量维度: " << record.descriptor.size() << std::endl;
    }
}

std::vector<dlib::rectangle> FaceRecognizer::detectFaces(
    const dlib::matrix<dlib::rgb_pixel>& img) {
    return detector_(img);
}

dlib::matrix<dlib::rgb_pixel> FaceRecognizer::alignFace(
    const dlib::matrix<dlib::rgb_pixel>& img, 
    const dlib::rectangle& face_rect) {
    
    auto shape = sp_(img, face_rect);
    dlib::matrix<dlib::rgb_pixel> face_chip;
    dlib::extract_image_chip(img, 
        dlib::get_face_chip_details(shape, 150, 0.25), 
        face_chip);
    return face_chip;
}

std::string FaceRecognizer::recognizeFace(
    const dlib::matrix<dlib::rgb_pixel>& face_chip) {
    
    // 提取特征向量
    std::vector<dlib::matrix<dlib::rgb_pixel>> faces_to_process;
    faces_to_process.push_back(face_chip);
    std::vector<dlib::matrix<float,0,1>> face_descriptors = net_(faces_to_process);
    dlib::matrix<float,0,1> descriptor = face_descriptors[0];
    
    // 在数据库中查找最接近的人脸
    double min_distance = threshold_ + 1.0; // 初始化为大于阈值
    std::string best_match = "陌生人";
    
    for (const auto& record : face_lib_) {
        double distance = dlib::length(descriptor - record.descriptor);
        if (distance < min_distance) {
            min_distance = distance;
            best_match = record.name;
        }
    }
    
    // 如果最小距离大于阈值，则认为是陌生人
    if (min_distance > threshold_) {
        return "Stranger";
    }
    
    return best_match;
}